import matplotlib
matplotlib.use('TkAgg')  # 使用 TkAgg 后端弹窗显示
import matplotlib.pyplot as plt
import numpy as np
import os

def plot_population_histograms(population_data, n_samples, distribution, output_path=None, low_age=0, up_age=100):
    """
    绘制给定人口数据的四张图：男性年龄分布直方图、女性年龄分布直方图、
    年龄-身高关系曲线图、年龄-体重关系曲线图，并可选保存到指定路径。

    参数:
        population_data: process_population_data 返回的元组
        n_samples: 总样本数（用于转换为频率百分比）
        distribution: 分布名称（用于标题）
        output_path: 保存图形的目录（可选）
        low_age: 年龄下限（用于限制 X 轴）
        up_age: 年龄上限（用于限制 X 轴）
    """
    (ages_male, ages_female), (heights_male, heights_female), \
    (weights_male, weights_female), (bsas_male, bsas_female), _ = population_data

    # 1. 男性年龄分布直方图
    fig_male_age = plt.figure(figsize=(8, 6))
    ax_male_age = fig_male_age.add_subplot(111)
    counts_male, _, _ = ax_male_age.hist(ages_male, bins=20, color='blue', alpha=0.5, label='Male')
    ax_male_age.set_xlabel('Age (years)', fontsize=11)
    ax_male_age.set_ylabel('Frequency (%)', fontsize=11)
    ax_male_age.legend(fontsize=9)
    ax_male_age.grid(True, linestyle='--', alpha=0.6)
    # 百分比 Y 轴
    max_count_male = np.max(counts_male)
    yticks_male = np.linspace(0, max_count_male, 6)
    ax_male_age.set_yticks(yticks_male)
    ax_male_age.set_yticklabels([f'{int(y / len(ages_male) * 100)}%' for y in yticks_male], fontsize=9)
    ax_male_age.tick_params(axis='x', labelsize=10)
    ax_male_age.set_ylim(top=max_count_male * 1.15)
    ax_male_age.set_xlim(low_age, up_age)  # 限制 X 轴
    ax_male_age.set_xticks(np.linspace(low_age, up_age, 9))  # 统一 X 轴刻度
    fig_male_age.subplots_adjust(left=0.12, right=0.95, top=0.9, bottom=0.15)
    fig_male_age.suptitle(f"{distribution} - Male Age Distribution", fontsize=14)
    if output_path:
        fig_male_age.savefig(os.path.join(output_path, f"{distribution}_male_age_histogram.png"))
    plt.show()

    # 2. 女性年龄分布直方图
    fig_female_age = plt.figure(figsize=(8, 6))
    ax_female_age = fig_female_age.add_subplot(111)
    counts_female, _, _ = ax_female_age.hist(ages_female, bins=20, color='red', alpha=0.5, label='Female')
    ax_female_age.set_xlabel('Age (years)', fontsize=11)
    ax_female_age.set_ylabel('Frequency (%)', fontsize=11)
    ax_female_age.legend(fontsize=9)
    ax_female_age.grid(True, linestyle='--', alpha=0.6)
    # 百分比 Y 轴
    max_count_female = np.max(counts_female)
    yticks_female = np.linspace(0, max_count_female, 6)
    ax_female_age.set_yticks(yticks_female)
    ax_female_age.set_yticklabels([f'{int(y / len(ages_female) * 100)}%' for y in yticks_female], fontsize=9)
    ax_female_age.tick_params(axis='x', labelsize=10)
    ax_female_age.set_ylim(top=max_count_female * 1.15)
    ax_female_age.set_xlim(low_age, up_age)  # 限制 X 轴
    ax_female_age.set_xticks(np.linspace(low_age, up_age, 9))  # 统一 X 轴刻度
    fig_female_age.subplots_adjust(left=0.12, right=0.95, top=0.9, bottom=0.15)
    fig_female_age.suptitle(f"{distribution} - Female Age Distribution", fontsize=14)
    if output_path:
        fig_female_age.savefig(os.path.join(output_path, f"{distribution}_female_age_histogram.png"))
    plt.show()

    # 3. 年龄-身高关系曲线图
    fig_height = plt.figure(figsize=(8, 6))
    ax_height = fig_height.add_subplot(111)
    ax_height.scatter(ages_male, heights_male, alpha=0.5, color='blue', label='Male', s=20)
    ax_height.scatter(ages_female, heights_female, alpha=0.5, color='red', label='Female', s=20)
    # 计算并绘制趋势线（使用多项式拟合）
    male_z = np.polyfit(ages_male, heights_male, 2)
    female_z = np.polyfit(ages_female, heights_female, 2)
    x_range = np.linspace(low_age, up_age, 100)
    male_poly = np.poly1d(male_z)
    female_poly = np.poly1d(female_z)
    ax_height.plot(x_range, male_poly(x_range), 'b-', linewidth=2)
    ax_height.plot(x_range, female_poly(x_range), 'r-', linewidth=2)
    ax_height.set_xlabel('Age (years)', fontsize=12)
    ax_height.set_ylabel('Height (cm)', fontsize=12)
    ax_height.set_title(f"{distribution} - Age vs Height Relationship", fontsize=14)
    ax_height.grid(True, linestyle='--', alpha=0.6)
    ax_height.legend(fontsize=10)
    ax_height.set_xlim(low_age, up_age)  # 限制 X 轴
    ax_height.set_xticks(np.linspace(low_age, up_age, 9))  # 统一 X 轴刻度
    fig_height.subplots_adjust(left=0.12, right=0.95, top=0.9, bottom=0.15)
    if output_path:
        fig_height.savefig(os.path.join(output_path, f"{distribution}_age_height_relationship.png"))
    plt.show()

    # 4. 年龄-体重关系曲线图
    fig_weight = plt.figure(figsize=(8, 6))
    ax_weight = fig_weight.add_subplot(111)
    ax_weight.scatter(ages_male, weights_male, alpha=0.5, color='blue', label='Male', s=20)
    ax_weight.scatter(ages_female, weights_female, alpha=0.5, color='red', label='Female', s=20)
    # 计算并绘制趋势线
    male_z = np.polyfit(ages_male, weights_male, 2)
    female_z = np.polyfit(ages_female, weights_female, 2)
    male_poly = np.poly1d(male_z)
    female_poly = np.poly1d(female_z)
    ax_weight.plot(x_range, male_poly(x_range), 'b-', linewidth=2)
    ax_weight.plot(x_range, female_poly(x_range), 'r-', linewidth=2)
    ax_weight.set_xlabel('Age (years)', fontsize=12)
    ax_weight.set_ylabel('Weight (kg)', fontsize=12)
    ax_weight.set_title(f"{distribution} - Age vs Weight Relationship", fontsize=14)
    ax_weight.grid(True, linestyle='--', alpha=0.6)
    ax_weight.legend(fontsize=10)
    ax_weight.set_xlim(low_age, up_age)  # 限制 X 轴
    ax_weight.set_xticks(np.linspace(low_age, up_age, 9))  # 统一 X 轴刻度
    fig_weight.subplots_adjust(left=0.12, right=0.95, top=0.9, bottom=0.15)
    if output_path:
        fig_weight.savefig(os.path.join(output_path, f"{distribution}_age_weight_relationship.png"))
    plt.show()